import os
import torch

from pathlib import Path
from PIL import Image


class Forecase:

    def __init__(self, model, transform):
        self.model = model
        self.transform = transform

    def run_forecast(self, filePath, classes):
        # 查看模型
        print(self.model.eval())
        image_dir = filePath
        # image_dir = '/Users/xianda/Downloads/flower7595/flowers/dandelion'
        index = 0
        for filename in Path(image_dir).glob('*.jpg'):
            try:
                file_path = os.path.join(image_dir, filename)

                # 读取输入图片
                image_path = file_path
                input_image = Image.open(image_path)
                input_image = self.transform(input_image)  # 预处理输入图片
                input_image = input_image.unsqueeze(0)  # 添加 batch 维度

                # 使用模型进行推理
                with torch.no_grad():
                    output = self.model(input_image)

                # 解释模型输出
                probabilities = torch.softmax(output, dim=1)
                max_prob, predicted_class = torch.max(probabilities, 1)

                index += 1
                # 打印结果
                print(f"正在进行第{index}个文件的预测：{file_path}")
                print("类别为:", classes[predicted_class.item()])
                print("准确率:", max_prob.item())
            except Exception as e:
                print(f'File is {os.path.join(image_dir, filename)} Exception:{e}')
                continue